package com.shawn.datastructure.trie;

/**
 * 字典树
 *
 * @author shawn
 */
public class TrieTree {

	/**
	 * 前缀树的结构
	 */
	private static class TrieNode {

		private int path;

		private int end;

		private TrieNode[] nexts;

		public TrieNode() {
			path = 0;
			end = 0;
			nexts = new TrieNode[26];
		}
	}

	/**
	 * 前缀树的头结点 为空
	 */
	private TrieNode root;

	public TrieTree() {
		root = new TrieNode();
	}

	public static void main(String[] args) {
		TrieTree trie = new TrieTree();
		trie.insert("abcd");
		trie.insert("efd");
		trie.insert("ggg");
		trie.insert("rrr");
		trie.insert("ggg");
		trie.insert("gff");
		trie.delete("rrr");
		System.out.println(trie.search(("g")));
		System.out.println(trie.prefixNumber(("gg")));
	}

	/**
	 * 将一个字符串加入前缀树，即向前缀树中添加路径和结点，
	 * <p>
	 * 首先检查前缀树中是否含有当前结点，
	 * <p>
	 * 如果存在的话，则将其路径上的path++，end++，
	 * <p>
	 * 如果不存在某几个结点的话，则创建出来
	 *
	 * @param str
	 */
	public void insert(String str) {
		if (str == null || str.length() == 0) {
			return;
		}
		char[] ch = str.toCharArray();

		TrieNode trieNode = root;

		for (int i = 0; i < ch.length; i++) {
			int index = ch[i] - 'a';
			//该元素不存在，则创建出来
			if (trieNode.nexts[index] == null) {
				trieNode.nexts[index] = new TrieNode();
			}
			trieNode = trieNode.nexts[index];
			trieNode.path++;
		}
		trieNode.end++;
	}

	/**
	 * 删除指定字符串对应在前缀树中的路径 在前缀树中查找，然后依次将其path-1，
	 * <p>
	 * 当path==0时，直接将其置为null即可（后面的结点都可以不要了）
	 *
	 * @param str
	 */
	public void delete(String str) {
		if (str == null || str.length() == 0) {
			return;
		}
		char[] ch = str.toCharArray();
		TrieNode trieNode = root;

		for (int i = 0; i < ch.length; i++) {
			int index = ch[i] - 'a';
			if (--trieNode.nexts[index].path == 0) {
				trieNode.nexts[index] = null;
				return;
			}
			trieNode = trieNode.nexts[index];
		}
		trieNode.end--;
	}

	/**
	 * 查询一个字符串在前缀树中出现了多少次 如果一路下来都存在，
	 * <p>
	 * 则返回end中的值， 如果过程中含有不存在的值，则返回0
	 *
	 * @param str
	 *
	 * @return
	 */
	public int search(String str) {
		if (str == null || str.length() == 0) {
			return 0;
		}
		char[] ch = str.toCharArray();

		TrieNode trieNode = root;
		for (int i = 0; i < ch.length; i++) {
			int index = ch[i] - 'a';
			if (trieNode.nexts[index] == null) {
				return 0;
			}
			trieNode = trieNode.nexts[index];
		}
		return trieNode.end;
	}

	public int prefixNumber(String str) {
		if (str == null || str.length() == 0) {
			return 0;
		}
		char[] ch = str.toCharArray();
		TrieNode trieNode = root;

		for (int i = 0; i < ch.length; i++) {
			int index = ch[i] - 'a';
			if (trieNode.nexts[index] == null) {
				return 0;
			}
			trieNode = trieNode.nexts[index];
		}
		getString(trieNode);
		return trieNode.path;
	}

	public void getString(TrieNode trieNode) {
		TrieNode[] nodes = trieNode.nexts;
		for (int i = 0; i < nodes.length; i++) {
			if (nodes[i] != null) {
				getString(nodes[i]);
			}
		}
	}


}
